#include "tcp_sender.hh"

#include "tcp_config.hh"

#include <random>

// Dummy implementation of a TCP sender

// For Lab 3, please replace with a real implementation that passes the
// automated checks run by `make check_lab3`.

template <typename... Targs>
void DUMMY_CODE(Targs &&.../* unused */) {}

using namespace std;

//! \param[in] capacity the capacity of the outgoing byte stream
//! \param[in] retx_timeout the initial amount of time to wait before retransmitting the oldest outstanding segment
//! \param[in] fixed_isn the Initial Sequence Number to use, if set (otherwise uses a random ISN)
TCPSender::TCPSender(const size_t capacity, const uint16_t retx_timeout, const std::optional<WrappingInt32> fixed_isn)
    : _isn(fixed_isn.value_or(WrappingInt32{random_device()()}))
    , _initial_retransmission_timeout{retx_timeout}
    , _RTO(retx_timeout)
    , _stream(capacity) {}

uint64_t TCPSender::bytes_in_flight() const { return _bytes_in_flight; }

void TCPSender::fill_window() {
    if (_next_seqno == 0) {
        TCPSegment seg;
        seg.header().syn = true;
        send_seg(seg);
    }
    size_t win = _win_size > 0 ? _win_size : 1;
    size_t remain;
    if (win + _last_acked < _next_seqno)
        return;
    while ((remain = win - (_next_seqno - _last_acked)) != 0 && !_fin_flg) {
        size_t size = min(TCPConfig::MAX_PAYLOAD_SIZE, remain);
        TCPSegment seg;
        string str = _stream.read(size);
        seg.payload() = Buffer(std::move(str));
        // add FIN
        if (seg.length_in_sequence_space() < win && _stream.eof()) {
            seg.header().fin = _fin_flg = true;
        }
        // stream is empty
        if (seg.length_in_sequence_space() == 0)
            return;

        send_seg(seg);
    }
}

void TCPSender::send_seg(TCPSegment &seg) {
    auto &header = seg.header();
    header.seqno = next_seqno();
    header.syn = _next_seqno == 0;
    //    header.fin = _fin_flg;
    _next_seqno += seg.length_in_sequence_space();
    _bytes_in_flight += seg.length_in_sequence_space();

    _segments_out.push(seg);
    _outstanding_segs.push_back(seg);
    _timer_on = true;
}

//! \param ackno The remote receiver's ackno (acknowledgment number)
//! \param window_size The remote receiver's advertised window size
void TCPSender::ack_received(const WrappingInt32 ackno, const uint16_t window_size) {
    size_t abs_ackno = unwrap(ackno, _isn, _last_acked);
    _win_size = window_size;
    if (abs_ackno <= _last_acked || abs_ackno > _next_seqno)
        return;
    auto it = _outstanding_segs.begin();
    size_t acc = _last_acked;
    while (it not_eq _outstanding_segs.end() && acc < abs_ackno) {
        acc += it->length_in_sequence_space();
        ++it;
    }
//    if (acc < abs_ackno)
//        return;
    if (acc > abs_ackno) {
        --it;
        acc -= it->length_in_sequence_space();
    }
    if (it == _outstanding_segs.begin())
        return;
    _bytes_in_flight -= (acc - _last_acked);
    _outstanding_segs.erase(_outstanding_segs.begin(), it);
    _last_acked = acc;

    _timer_on = false;
    _RTO = _initial_retransmission_timeout;
    _consecutive_retransmissions = 0;
    _ticks = 0;

    fill_window();
    if (!_outstanding_segs.empty())
        _timer_on = true;
}

//! \param[in] ms_since_last_tick the number of milliseconds since the last call to this method
void TCPSender::tick(const size_t ms_since_last_tick) {
    if (not _timer_on)
        return;
    _ticks += ms_since_last_tick;
    if (_ticks >= _RTO && !_outstanding_segs.empty()) {
        _segments_out.push(_outstanding_segs.front());
        if (_win_size not_eq 0 || _last_acked == 0) {
            _consecutive_retransmissions += 1;
            _RTO <<= 1;
        }
        _ticks = 0;
    }
    if (_outstanding_segs.empty()) {
        _timer_on = false;
        _ticks = 0;
    }
}

unsigned int TCPSender::consecutive_retransmissions() const { return _consecutive_retransmissions; }

void TCPSender::send_empty_segment() {
    TCPSegment empty_segment;
    empty_segment.header().seqno = wrap(_next_seqno, _isn);
    _segments_out.push(empty_segment);
}
